import socket
import typing as tp

import aio_pika
from loguru import logger

from ..agent import K2CIAgent


class RabbitMQClient:
    MAX_RETRY_CONNECT: int = 5

    def __init__(self, agent: K2CIAgent):
        self.agent = agent
        self.connection: tp.Optional[aio_pika.Connection] = None
        self.channel: tp.Optional[aio_pika.Channel] = None

    async def connect(self) -> bool:
        result = False
        try:
            self.connection = await aio_pika.connect_robust(
                f"amqp://{self.agent.cfg.rabbit_username}:{self.agent.cfg.rabbit_passwd}"
                f"@{self.agent.cfg.rabbit_mq_ip}:{self.agent.cfg.rabbit_port}",
                timeout=5
            )
            self.channel = await self.connection.channel()
            logger.info(f"[K2CI] amqp://{self.agent.cfg.rabbit_mq_ip}:{self.agent.cfg.rabbit_port} 连接成功")
            result = True
        except (aio_pika.exceptions.AMQPConnectionError, ConnectionError, TimeoutError) as e:
            logger.error(f"[K2CI] amqp://{self.agent.cfg.rabbit_mq_ip}:{self.agent.cfg.rabbit_port} 连接失败：{e}")
        except (BaseException, Exception):
            import traceback
            logger.error(traceback.format_exc())

        return result

    async def try_connect(self, ignore_max_retry: bool = False, ignore_conn_error: bool = False) -> bool:
        if self.is_online:
            return True

        if ignore_max_retry:
            while True:
                if await self.connect():
                    return True
        else:
            for i in range(self.MAX_RETRY_CONNECT):
                if await self.connect():
                    return True
            else:
                if not ignore_conn_error:
                    raise RuntimeError(f"[K2CI] amqp://{self.agent.cfg.rabbit_mq_ip}:{self.agent.cfg.rabbit_port} 连接失败")
                return False

    async def close(self):
        if self.connection and not self.connection.is_closed:
            await self.connection.close()

    @property
    def is_online(self):
        return self.connection and not self.connection.is_closed
